package com.zzl.study.auth.security.service;

import com.zzl.study.auth.security.authentication.client.GatewayOAuth2AccessToken;
import com.zzl.study.auth.security.authentication.client.GatewayOAuth2AuthorizedClient;
import com.zzl.study.auth.security.authentication.client.GatewayOAuth2RefreshToken;
import org.springframework.data.redis.core.ReactiveRedisTemplate;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.client.OAuth2AuthorizedClient;
import org.springframework.security.oauth2.client.OAuth2AuthorizedClientId;
import org.springframework.security.oauth2.client.ReactiveOAuth2AuthorizedClientService;
import org.springframework.security.oauth2.client.registration.ReactiveClientRegistrationRepository;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2RefreshToken;
import org.springframework.util.Assert;
import reactor.core.publisher.Mono;

/**
 * @author: zhangzl
 * @date: 2024/6/11 14:11
 * @version: 1.0
 * @description: 基于Redis的授权客户端服务
 */
public class RedisReactiveOAuth2AuthorizedClientService implements ReactiveOAuth2AuthorizedClientService {

    private final ReactiveClientRegistrationRepository clientRegistrationRepository;

    private final ReactiveRedisTemplate<String, GatewayOAuth2AuthorizedClient> oauth2AuthorizedClientRedisTemplate;

    private static final String DEFAULT_AUTHORIZED_CLIENT_KEY_PREFIX = "OAUTH2_AUTHORIZED_CLIENT";

    private static final String DEFAULT_AUTHORIZED_CLIENT_KEY_FORMAT = "%s:%s-%s";

    public RedisReactiveOAuth2AuthorizedClientService(ReactiveClientRegistrationRepository clientRegistrationRepository,
                                                      ReactiveRedisTemplate<String, GatewayOAuth2AuthorizedClient> oauth2AuthorizedClientRedisTemplate) {
        this.clientRegistrationRepository = clientRegistrationRepository;
        this.oauth2AuthorizedClientRedisTemplate = oauth2AuthorizedClientRedisTemplate;
    }

    @Override
    @SuppressWarnings("unchecked")
    public <T extends OAuth2AuthorizedClient> Mono<T> loadAuthorizedClient(String clientRegistrationId, String principalName) {
        Assert.hasText(clientRegistrationId, "clientRegistrationId cannot be empty");
        Assert.hasText(principalName, "principalName cannot be empty");
        String oAuth2AuthorizedClientCacheKey = getOAuth2AuthorizedClientCacheKey(clientRegistrationId, principalName);
        return (Mono<T>) this.oauth2AuthorizedClientRedisTemplate.opsForValue().get(oAuth2AuthorizedClientCacheKey)
                .flatMap((oAuth2AuthorizedClient) -> this.clientRegistrationRepository.findByRegistrationId(clientRegistrationId)
                        .flatMap((clientRegistration) -> {
                            OAuth2AccessToken oAuth2AccessToken = GatewayOAuth2AccessToken.to(oAuth2AuthorizedClient.getAccessToken());
                            OAuth2RefreshToken oAuth2RefreshToken = GatewayOAuth2RefreshToken.to(oAuth2AuthorizedClient.getRefreshToken());
                            return Mono.just(new OAuth2AuthorizedClient(clientRegistration, principalName, oAuth2AccessToken, oAuth2RefreshToken));
                        }));
    }

    @Override
    public Mono<Void> saveAuthorizedClient(OAuth2AuthorizedClient authorizedClient, Authentication principal) {
        Assert.notNull(authorizedClient, "authorizedClient cannot be null");
        Assert.notNull(principal, "principal cannot be null");
        GatewayOAuth2AuthorizedClient gatewayOAuth2AuthorizedClient = GatewayOAuth2AuthorizedClient.of(authorizedClient);
        String oAuth2AuthorizedClientCacheKey = getOAuth2AuthorizedClientCacheKey(authorizedClient.getClientRegistration().getRegistrationId(), principal.getName());
        return this.oauth2AuthorizedClientRedisTemplate.opsForValue().set(oAuth2AuthorizedClientCacheKey, gatewayOAuth2AuthorizedClient)
                .flatMap(success -> {
                    if (success) {
                        return Mono.empty();
                    } else {
                        return Mono.error(new IllegalStateException("Failed to save OAuth2AuthorizedClient"));
                    }
                });
    }

    @Override
    public Mono<Void> removeAuthorizedClient(String clientRegistrationId, String principalName) {
        Assert.hasText(clientRegistrationId, "clientRegistrationId cannot be empty");
        Assert.hasText(principalName, "principalName cannot be empty");
        String oAuth2AuthorizedClientCacheKey = getOAuth2AuthorizedClientCacheKey(clientRegistrationId, principalName);
        return this.oauth2AuthorizedClientRedisTemplate.opsForValue().delete(oAuth2AuthorizedClientCacheKey)
                .flatMap(success -> {
                    if (success) {
                        return Mono.empty();
                    } else {
                        return Mono.error(new IllegalStateException("Failed to remove OAuth2AuthorizedClient"));
                    }
                });
    }

    private static String getOAuth2AuthorizedClientCacheKey(String clientRegistrationId, String principalName) {
        return String.format(DEFAULT_AUTHORIZED_CLIENT_KEY_FORMAT, DEFAULT_AUTHORIZED_CLIENT_KEY_PREFIX, clientRegistrationId, principalName);
    }

}
